'''
Vive Ultimate Tracker Interface.

Author: Jin Liu.
Date: 2025/06/16
'''
import time
import openvr
from ezgl import Matrix4x4
# from line_profiler import LineProfilerc:\Users\LENOVO\Desktop\DexCap\STEP1_collect_data_202408updates\configs\OG000086 c:\Users\LENOVO\Desktop\DexCap\STEP1_collect_data_202408updates\configs\OG000131


class ViveTracker:
    def __init__(self, index=1):
        # 初始化OpenVR
        self.vr_system = None
        self.trackers = {}
        self.init_openvr()
        self.index = index
        if index not in self.trackers:
            raise ValueError(f"Device index {index} not found in trackers.")

    def init_openvr(self):
        """初始化OpenVR并查找所有Vive Tracker设备"""
        try:
            print("OpenVR init...")
            openvr.init(openvr.VRApplication_Other)
        except Exception as e:
            print(f"OpenVR init failed: {e}")
            print(f"Retrying OpenVR init...")
            openvr.shutdown()
            time.sleep(0.5)
            openvr.init(openvr.VRApplication_Other)

        self.vr_system = openvr.VRSystem()
        self.find_trackers()

    def find_trackers(self):
        """查找并识别所有Vive Tracker设备"""
        for i in range(openvr.k_unMaxTrackedDeviceCount):
            device_class = self.vr_system.getTrackedDeviceClass(i)
            
            if device_class == openvr.TrackedDeviceClass_GenericTracker:
                serial = self.vr_system.getStringTrackedDeviceProperty(i, openvr.Prop_SerialNumber_String)
                device_name = self.vr_system.getStringTrackedDeviceProperty(i, openvr.Prop_RenderModelName_String)
                self.trackers[i] = {
                    'serial': serial,
                    'name': device_name,
                    'index': i
                }
                print(f"Found Vive Tracker: index={i}, serial={serial}, name={device_name}")
    
    # @profile
    def recv_data(self):
        """获取指定索引的Vive Tracker的姿态数据"""
        # 获取设备的姿态
        data = self.vr_system.getDeviceToAbsoluteTrackingPose(
            openvr.TrackingUniverseStanding, 0, openvr.k_unMaxTrackedDeviceCount
        )
        if data[self.index].bPoseIsValid:
            mat = data[self.index].mDeviceToAbsoluteTracking
            pose = Matrix4x4([
                [mat[0][0], mat[0][1], mat[0][2], mat[0][3]],
                [mat[1][0], mat[1][1], mat[1][2], mat[1][3]],
                [mat[2][0], mat[2][1], mat[2][2], mat[2][3]],
                [0, 0, 0, 1]
            ])
            # time.sleep(0.01)
            return pose
        return None
    
    def close(self):
        """关闭OpenVR系统"""
        if self.vr_system:
            openvr.shutdown()
            print("OpenVR Closed")



if __name__ == "__main__":
    from ezrobot.sensor.base import SensorManager, accurate_sleep

    tracker = ViveTracker(1)
    sensor = SensorManager(sensor=tracker, frequency=30, buffer_size=10, streaming=True)
    try:
        while True:
            accurate_sleep(0.1)
            print(sensor.get_data())
    except Exception as e:
        print(e)
        print("Stopping encoder streaming...")
        sensor.stop_streaming()